"""
一个简单的demo，调用CharacterGLM实现角色扮演，调用CogView生成图片，调用ChatGLM生成CogView所需的prompt。

依赖：
pyjwt
requests
streamlit
zhipuai
python-dotenv

运行方式：
```bash
streamlit run characterglm_api_demo_streamlit.py
```
"""
import os
import itertools
from typing import Iterator, Optional

import streamlit as st
from dotenv import load_dotenv

# 通过.env文件设置环境变量
# reference: https://github.com/theskumar/python-dotenv
load_dotenv()

import api
from api import get_characterglm_response
from data_types import TextMsg, filter_text_msg

st.set_page_config(page_title="CharacterGLM API Demo", page_icon="🤖", layout="wide")
debug = os.getenv("DEBUG", "").lower() in ("1", "yes", "y", "true", "t", "on")


def update_api_key(key: Optional[str] = None):
    if debug:
        print(f'update_api_key. st.session_state["API_KEY"] = {st.session_state["API_KEY"]}, key = {key}')
    key = key or st.session_state["API_KEY"]
    if key:
        api.API_KEY = key


# 设置API KEY
api_key = st.sidebar.text_input("API_KEY", value=os.getenv("API_KEY", ""), key="API_KEY", type="password",
                                on_change=update_api_key)
update_api_key(api_key)

# 初始化
if "history" not in st.session_state:
    st.session_state["history"] = []
if "meta" not in st.session_state:
    st.session_state["meta"] = {
        "bot_info": "小白，性别女，17岁，平溪孤儿院的孩子。小白患有先天的白血病，头发为银白色。小白身高158cm，体重43kg。小白的名字是孤儿院院长给起的名字，因为小白是在漫天大雪白茫茫的一片土地上被捡到的。小白经常穿一身破旧的红裙子，只是为了让自己的气色看上去红润一些。小白初中毕业，没有上高中，学历水平比较低。小白在孤儿院相处最好的一个人是阿南，小白喊阿南哥哥。阿南对小白很好。",
        "user_info": "阿南是小白的哥哥",
        "bot_name": "小白",
        "user_name": "阿南"
    }


def init_session():
    st.session_state["history"] = []


# 4个输入框，设置meta的4个字段
meta_labels = {
    "bot_name": "角色1名字",
    "user_name": "角色2名字",
    "bot_info": "角色1人设",
    "user_info": "角色2人设"
}

# 2x2 layout
with st.container():
    col1, col2 = st.columns(2)
    with col1:
        st.text_input(label="角色1", value=st.session_state["meta"]["bot_name"], key="bot_name",
                      on_change=lambda: st.session_state["meta"].update(role1_name=st.session_state["bot_name"]),
                      help="模型所扮演的角色的名字，不可以为空")
        st.text_area(label="角色1人设", value=st.session_state["meta"]["bot_info"], key="bot_info",
                     on_change=lambda: st.session_state["meta"].update(role1_info=st.session_state["bot_info"]),
                     help="角色的详细人设信息，不可以为空")

    with col2:
        st.text_input(label="角色2", value=st.session_state["meta"]["user_name"], key="user_name",
                      on_change=lambda: st.session_state["meta"].update(role2_name=st.session_state["user_name"]),
                      help="用户的名字，默认为用户")
        st.text_area(label="角色2人设", value=st.session_state["meta"]["user_info"], key="user_info",
                     on_change=lambda: st.session_state["meta"].update(role2_info=st.session_state["user_info"]),
                     help="用户的详细人设信息，不可以为空")


def verify_meta() -> bool:
    # 检查`角色名`和`角色人设`是否空，若为空，则弹出提醒
    if (st.session_state["meta"]["bot_name"] == ""
            or st.session_state["meta"]["bot_info"] == ""
            or st.session_state["meta"]["user_name"] == ""
            or st.session_state["meta"]["user_info"] == ""):
        st.error("角色名和角色人设不能为空")
        return False
    else:
        return True


button_labels = {
    "clear_meta": "清空人设",
    "clear_history": "清空对话历史"
}
if debug:
    button_labels.update({
        "show_api_key": "查看API_KEY",
        "show_meta": "查看meta",
        "show_history": "查看历史"
    })

# 在同一行排列按钮
with st.container():
    n_button = len(button_labels)
    cols = st.columns(n_button)
    button_key_to_col = dict(zip(button_labels.keys(), cols))

    with button_key_to_col["clear_meta"]:
        clear_meta = st.button(button_labels["clear_meta"], key="clear_meta")
        if clear_meta:
            st.session_state["meta"] = {
                "bot_info": "",
                "user_info": "",
                "bot_name": "",
                "user_name": ""
            }
            st.rerun()

    with button_key_to_col["clear_history"]:
        clear_history = st.button(button_labels["clear_history"], key="clear_history")
        if clear_history:
            init_session()
            st.rerun()

with st.chat_message(name="user", avatar="user"):
    role1_placeholder = st.empty()
with st.chat_message(name="assistant", avatar="assistant"):
    role2_placeholder = st.empty()


def output_stream_response(response_stream: Iterator[str], placeholder):
    content = ""
    for content in itertools.accumulate(response_stream):
        placeholder.markdown(content)
    return content


def start_chat():
    hint = "早上好啊，小白（微笑）"
    role = "user"
    chat_content = []
    if not hint:
        return

    for i in range(20):
        if not verify_meta():
            return
        if not api.API_KEY:
            st.error("未设置API_KEY")

        if role == "user":
            role1_placeholder.markdown(hint)
            placeholder = role2_placeholder
        else:
            role2_placeholder.markdown(hint)
            placeholder = role1_placeholder
        st.session_state["history"].append(TextMsg({"role": role, "content": hint}))
        response_stream = get_characterglm_response(filter_text_msg(st.session_state["history"]),
                                                    meta=st.session_state["meta"])
        bot_response = output_stream_response(response_stream, placeholder)

        if not bot_response:
            role2_placeholder.markdown("生成出错")
            st.session_state["history"].pop()

        if role == "user":
            role = "assistant"
        else:
            role = "user"
        chat_content.append(st.session_state["meta"]["user_name"] + ":" + hint)
        chat_content.append(st.session_state["meta"]["bot_name"] + ":" + bot_response)

        # TODO 虽然是自动生成对话，但是要保持人设，但是这里有点问题，到后面就乱了
        st.session_state["meta"]["bot_name"], st.session_state["meta"]["user_name"] = st.session_state["meta"][
            "user_name"], st.session_state["meta"]["bot_name"]
        st.session_state["meta"]["bot_info"], st.session_state["meta"]["user_info"] = st.session_state["meta"][
            "user_info"], st.session_state["meta"]["bot_info"]
        hint = bot_response

    # 把每一次的对话内容保存到文件里
    final_content = list(set(chat_content))
    final_content.sort(key=chat_content.index)
    with open('./dialog_content.txt', 'a', encoding='utf-8') as file:
        file.writelines("%s\n" % s for s in final_content)


start_chat()
